# src/algorithms/depth_estimation.py
import numpy as np
import cv2


class DepthEstimator:
    def __init__(self, disparity: np.ndarray, reprojection_matrix: np.ndarray):
        """初始化深度估计器

        Args:
            disparity (np.ndarray): 视差图, 单通道浮点型数组
            reprojection_matrix (np.ndarray): 重投影矩阵 Q, 4x4 浮点型数组
        """
        self._disparity = disparity
        self._Q = reprojection_matrix
        self._focal_length = self._Q[2, 3]  # 焦距 fx
        self._baseline = -1 / self._Q[3, 2]  # 基线长度 B

    @staticmethod
    def writePly(fn, verts, colors):
        ply_header = """ply
        format ascii 1.0
        element vertex %(vert_num)d
        property float x
        property float y
        property float z
        property uchar red
        property uchar green
        property uchar blue
        end_header
        """

        verts = verts.reshape(-1, 3)  # Nx3 浮点数组, 每行一条 3D 坐标
        colors = colors.reshape(-1, 3)  # 与 verts 同行的 Nx3 数组, 0-255 的 RGB
        verts = np.hstack(
            [verts, colors]
        )  #  # 水平拼接 → verts[i] = [x, y, z, r, g, b]

        with open(fn, "wb") as f:
            f.write((ply_header % dict(vert_num=len(verts))).encode("utf-8"))
            np.savetxt(f, verts, fmt="%f %f %f %d %d %d ")

    @staticmethod
    def disparity2Pointcloud(
        disparity: np.ndarray,
        reprojection_matrix: np.ndarray,
        left_image: np.ndarray,
        output_filename="out.ply",
        min_disp: float = 1.0,
        max_disp: float = 160.0,
    ):
        """将视差图转换为点云

        Args:
            disparity (np.ndarray): 视差图, 单通道浮点型数组
            reprojection_matrix (np.ndarray): Q 矩阵
            left_image (np.ndarray): 左图 (BGR)
            output_filename (str): 输出点云文件名
            min_disp (float): 最小有效视差 (过滤掉小于该值的点)
            max_disp (float): 最大有效视差 (过滤掉大于该值的点)
        """
        # 有效视差过滤
        mask = np.isfinite(disparity) & (disparity > min_disp) & (disparity < max_disp)

        # 重投影到 3D
        points = cv2.reprojectImageTo3D(
            disparity, reprojection_matrix, handleMissingValues=True
        )

        # 提取颜色
        colors = cv2.cvtColor(left_image, cv2.COLOR_BGR2RGB)

        # 应用 mask
        output_points = points[mask]
        output_colors = colors[mask]

        # 写入 ply 文件
        DepthEstimator.writePly(output_filename, output_points, output_colors)

        print(f"Point cloud saved to {output_filename}")

    def disparity2Depth(
        self,
        disparity: np.ndarray,
        min_depth: float = 0.1,
        max_depth: float = 100.0,
    ) -> np.ndarray:
        """将视差图转换为深度图 (Z 值)

        Z = f * B / disparity

        Args:
            disparity (np.ndarray): 视差图, 单通道浮点型数组
            min_depth (float): 最小深度 (米/毫米, 取决于标定单位)
            max_depth (float): 最大深度

        Returns:
            np.ndarray: 深度图 (与输入视差图同尺寸)，无效点为 np.nan
        """
        depth = np.full_like(disparity, np.nan, dtype=np.float32)

        valid_mask = np.isfinite(disparity) & (disparity > 0)
        depth[valid_mask] = (self._focal_length * self._baseline) / disparity[
            valid_mask
        ]

        # 过滤不在范围内的点
        depth[(depth < min_depth) | (depth > max_depth)] = np.nan

        return depth
